import joblib
from PyQt5 import QtCore, QtGui, QtWidgets
from PyQt5.QtWidgets import *
from PyQt5.QtCore import pyqtSignal,Qt
from PyQt5 import QtWidgets
from PyQt5.QtWidgets import QMessageBox
from PyQt5.QtGui import QStandardItemModel, QStandardItem
import MyWindow
import torch.nn as nn
import numpy as np
import pandas as pd
import torch
import pickle
from sklearn.linear_model import LogisticRegression
import os
from MyWindow import Ui_MainWindow
import sys

class MyApp(QMainWindow, Ui_MainWindow):
    def __init__(self):
        super().__init__()
        self.setupUi(self)  # 初始化窗口和其中的控件
       # self.table.horizontalHeader().setSectionResizeMode(QHeaderView.Stretch)
       # self.table.horizontalHeader().setSectionResizeMode(0, QHeaderView.ResizeToContents)#列宽自动分配
        self.table.resizeColumnsToContents()

        self.getyb.clicked.connect(self.getyb_clicked)
        self.start.clicked.connect(self.start_clicked)
        self.undo.clicked.connect(self.undo_clicked)

    def hanggao(self):
        # 获取默认行高
        default_row_height = 35

        # 设置所有行的高度为默认行高的1.5倍，但转换为整数（向下取整）
        new_row_height = int(default_row_height * 1.5)
        for row in range(self.table.rowCount()):
            self.table.setRowHeight(row, new_row_height)

    def getyb_clicked(self):

        # 获取表格的列数
        column_count = self.table.columnCount()

        # 遍历每一行
        for row in range(self.table.rowCount()):
            # 获取当前行的最后一列单元格的项
            item = self.table.item(row, column_count - 1)

            # 如果该项存在，则清除其内容
            if item:
                item.setText('')
                # 设置参数
        n = 100  # 你想从CSV中随机读取的行数
        input_file = 'total_small.csv'
        output_file = 'random_sample.csv'
        chunksize = 1000  # 读取CSV的块大小，可以根据内存调整

        # 初始化一个空的列表来存储随机选择的行和原始行号
        rows_selected = []
        original_indices = []

        # 全局行号计数器，从1开始（因为没有标题行）
        global_row_count = 1

        # 逐块读取CSV文件
        for chunk in pd.read_csv(input_file, chunksize=chunksize, header=None):
            # 计算当前块中可以选择的行数
            num_to_select = min(n - len(rows_selected), len(chunk))
            if num_to_select == 0:
                break

                # 从当前块中随机选择行
            random_indices = np.random.choice(len(chunk), num_to_select, replace=False)

            # 将随机选择的行和它们的原始行号（基于全局行号计数器）添加到列表中
            for idx in random_indices:
                rows_selected.append(chunk.iloc[idx].tolist())
                original_indices.append(global_row_count + idx)  # idx是块内的索引，需要加上global_row_count

            # 更新全局行号计数器
            global_row_count += len(chunk)

            # 将列表转换为DataFrame
        random_df = pd.DataFrame(rows_selected)  # 没有列名，所以这里不指定columns

        # 确保DataFrame中恰好有n条记录，如果有多余的，则删除它们
        if len(random_df) > n:
            random_df = random_df.head(n)

            # 将DataFrame写入CSV文件（没有列名）
        random_df.to_csv(output_file, index=False, header=False)

        # 输出被收录的记录的原始行号

        #original_indices = original_indices[1:]
        print("Original indices of selected rows:", original_indices)
        print(len(original_indices))

        """# 如果需要，重新读取CSV文件到df变量中
        df = pd.read_csv(output_file)

        # 选择您想要的列（注意：pandas中的索引是从0开始的）
        selected_columns = df.iloc[:, [0, 1, 14, 40, 52, 70]]

        # 假设self.table是一个QTableWidget的实例
        self.table.setRowCount(len(selected_columns))
        for row_index, row in selected_columns.iterrows():
            for col_index, value in enumerate(row):
                item = QtWidgets.QTableWidgetItem(str(value))
                self.table.setItem(row_index, col_index, item)

        """
        #显示到表格
        random_sample_csv = 'random_sample.csv'
        df = pd.read_csv(random_sample_csv,header=None)

        # 选择您想要的列（注意：pandas中的索引是从0开始的）
        selected_columns = df.iloc[:, [0, 1, 14, 40, 52, 70]]

        # 假设self.table是一个QTableWidget的实例
        # 设置QTableWidget的行数为原始DataFrame的行数
        self.table.setRowCount(len(df))  # 使用df的行数，而不是selected_columns的列数

        # 设置QTableWidget的列数为所选列的数量
        self.table.setColumnCount(len(selected_columns.columns)+1)

        # 填充QTableWidget
        for row_index, row in selected_columns.iterrows():
            for col_index, value in enumerate(row):
                item = QtWidgets.QTableWidgetItem(str(value))  # 确保括号正确
                self.table.setItem(row_index, col_index, item)

        # 调用self.hanggao()，假设它是更新或显示表格的某个方法
        self.hanggao()
        #去掉重复的第一行

        # 读取Normalization_total_extend.csv，注意这里header=None表示不包含列名
        normalized_df = pd.read_csv('Normalization_total_small.csv', header=None)

        # 使用original_indices选择对应的行
        selected_rows = normalized_df.iloc[original_indices]

        # 保存选定的行到新的CSV文件，包括列名（如果知道列名的话）
        # 如果不知道列名或者不想包含列名，则直接保存
        selected_rows.to_csv('random_sample_scaled.csv', index=False, header=None)  # 不包含列名
        print("数据加入ok")

    def start_clicked(self):
        # 检查行数是否为零
        if self.table.rowCount() == 0:
            QMessageBox.warning(self, "警告", "请先加载数据!")
            return
        # 检查是否选择模型
        if not self.modelgroup.checkedButton():
            QMessageBox.warning(self, "警告", "未选择模型!")
            return



        # 检查哪个单选按钮被选中
        Label = ["BENIGN","Bot", "DDoS", "DoS GoldenEye","DoS Hulk", "DoS Slowhttptest", "DoS slowloris","FTP-Patator",
                 "Heartbleed","Infiltration","PortScan",
                 "SSH-Patator","Web Attack",
                  "Web Attack", "Web Attack"]
        if self.ICNN.isChecked():
                print("ICNN")

                class CNN(nn.Module):
                    def __init__(self):
                        super().__init__()
                        self.backbone = nn.Sequential(
                            nn.Conv1d(1, 32, kernel_size=2),  # 一个一维卷积层，输入通道数为1，输出通道数为32，卷积核大小为2
                            nn.Conv1d(32, 64, kernel_size=2),  # 另一个一维卷积层，输入通道数为32（与前一个卷积层的输出通道数匹配），输出通道数为64，卷积核大小为2
                            nn.MaxPool1d(2, 2),  # 一个一维最大池化层，池化窗口大小为2，步长也为2
                            nn.Conv1d(64, 64, kernel_size=2),  # 一个一维卷积层，输入通道数为64，输出通道数仍为64，卷积核大小为2
                            nn.Conv1d(64, 128, kernel_size=2),  # 另一个一维卷积层，输入通道数为64，输出通道数为128，卷积核大小为2
                            nn.MaxPool1d(2, 2),  # 另一个一维最大池化层
                        )
                        self.flatten = nn.Flatten()  # self.flatten是一个nn.Flatten层，用于将多维的输入一维化，以便可以将其输入到全连接层中
                        self.fc = nn.Sequential(  # self.fc是另一个nn.Sequential容器，它包含了多个全连接层（线性层）
                            nn.Linear(2304, 64),  # 一个全连接层，输入特征数为2304（这个数字通常是根据backbone的输出形状计算得出的），输出特征数为64
                            nn.ReLU(),  # 一个ReLU激活函数
                            nn.Linear(64, 64),  # 另一个全连接层，输入和输出特征数都是64
                            nn.ReLU(),  # 又一个ReLU激活函数
                            nn.Linear(64, 15)  # 最后一个全连接层，输入特征数为64，输出特征数为y_dimension（类别数）
                        )

                    def forward(self, X):
                        X = self.backbone(X)
                        X = self.flatten(X)
                        logits = self.fc(X)
                        return logits

                df = pd.read_csv('random_sample_scaled.csv',header=None)
                X = df.iloc[:, :-1].values  # 特征（不包含标签列）
                y = df.iloc[:, -1].values  # 标签（可选，这里仅用于示例）

                # 数据预处理：转换为浮点数张量，并添加通道维度
                X_tensor = torch.from_numpy(X).float().unsqueeze(1)  # 添加通道维度

                # 检查设备
                device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

                # 实例化模型
                model = CNN()

                # 加载预训练权重
                model.load_state_dict(torch.load('ICNNModel/CNN_model.pth', map_location=device))

                # 将模型和数据移动到GPU（如果可用）
                model = model.to(device)
                X_tensor = X_tensor.to(device)

                # 设置模型为评估模式
                model.eval()

                # 关闭梯度计算
                with torch.no_grad():
                    # 进行前向传播
                    outputs = model(X_tensor)
                # （可选）将logits转换为概率或类标签
                probabilities = torch.nn.functional.softmax(outputs, dim=1)
                predictions = torch.argmax(outputs, dim=1)


                # 将预测索引转换为标签
                predicted_labels = [Label[idx] for idx in predictions.tolist()]
                print(predictions)
                print(predicted_labels)


        elif self.LSTM.isChecked():
                print("LSTM")

                class LSTMModel(nn.Module):
                    def __init__(self, input_dim=78, hidden_dim=64, output_dim=15, num_layers=2, dropout=0.2):
                        super(LSTMModel, self).__init__()
                        self.hidden_dim = hidden_dim
                        self.num_layers = num_layers
                        self.lstm = nn.LSTM(input_dim, hidden_dim, num_layers, batch_first=True, dropout=dropout)
                        self.dropout = nn.Dropout(dropout)
                        self.fc = nn.Linear(hidden_dim, output_dim)

                    def forward(self, x):
                        h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_dim).to(device)
                        c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_dim).to(device)

                        out, _ = self.lstm(x, (h0, c0))
                        out = self.dropout(out)
                        out = self.fc(out[:, -1, :])  # 使用最后一个时间步的输出
                        return out

                df = pd.read_csv('random_sample_scaled.csv', header=None)
                X = df.iloc[:, :-1].values  # 特征（不包含标签列）
                y = df.iloc[:, -1].values  # 标签（可选，这里仅用于示例）

                # 数据预处理：转换为浮点数张量，并添加通道维度
                X_tensor = torch.from_numpy(X).float().unsqueeze(1)  # 添加通道维度

                # 检查设备
                device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

                # 实例化模型
                model = LSTMModel()

                # 加载预训练权重
                model.load_state_dict(torch.load('LSTMModel/LSTM_model.pth', map_location=device))

                # 将模型和数据移动到GPU（如果可用）
                model = model.to(device)
                X_tensor = X_tensor.to(device)

                # 设置模型为评估模式
                model.eval()

                # 关闭梯度计算
                with torch.no_grad():
                    # 进行前向传播
                    outputs = model(X_tensor)
                # （可选）将logits转换为概率或类标签
                probabilities = torch.nn.functional.softmax(outputs, dim=1)
                predictions = torch.argmax(outputs, dim=1)

                # 将预测索引转换为标签
                predicted_labels = [Label[idx] for idx in predictions.tolist()]
                print(predictions)
                print(predicted_labels)


        elif self.RF.isChecked():
                print("RF")
                # 1. 加载随机森林模型
                model_path = 'RandomForestModel/random_forest_model.pkl'
                rf_model = joblib.load(model_path)
                # 2. 加载数据集
                data_path = 'random_sample_scaled.csv'
                df = pd.read_csv(data_path,header=None)
                # 3. 假设数据集的最后一列是标签'Label'，我们先将其与特征分离
                X = df.iloc[:, :-1]  # 所有的行，除了最后一列的所有列（即特征）
                y_true = df.iloc[:, -1]  # 所有的行，最后一列（即真实标签）
                # 4. 使用模型进行预测
                y_pred = rf_model.predict(X)

                print(y_pred)
                predicted_labels = [Label[idx] for idx in y_pred.tolist()]
                print(predicted_labels)


        elif self.LOG.isChecked():
                print("LOG")
                model_file = 'LogisticRegressionModel/logistic_regression_model.pkl'
                logistic_regression_model = joblib.load(model_file)

                # 2. 读取CSV文件并处理数据
                data = pd.read_csv('random_sample_scaled.csv',header=None)
                # 假设最后一列是你要预测的目标列（通常在实际预测时不会包含），这里我们仅选择特征列
                X = data.iloc[:, :-1]  # 所有行，除了最后一列
                y_true = data.iloc[:, -1]  # 所有行，最后一列（真实标签，用于后续验证或比较）
                y_pred = logistic_regression_model.predict(X)

                # 4. 转换预测结果为标签
                predicted_labels = [Label[idx] for idx in y_pred]

                # 5. 输出结果
                print(y_pred)  # 输出预测结果的索引
                print(predicted_labels)  # 输出预测结果的标签


        #结果显示到表格
        row_count = self.table.rowCount()

        # 确保predicted_labels的长度与行数匹配
        #print(len(predicted_labels))
        #print(row_count)
        assert len(predicted_labels) == row_count, "predicted_labels的长度必须与table的行数一致"

        # 遍历每一行，将predicted_labels中的标签添加到最后一列
        for row_index in range(row_count):
            label_item = QTableWidgetItem(str(predicted_labels[row_index]))
            # 将label_item添加到当前行的最后一列（即列索引为列数减1的位置）
            column_index = self.table.columnCount() - 1
            self.table.setItem(row_index, column_index, label_item)





    def undo_clicked(self):
        # 清除表格中的所有数据，但保留表头和行/列结构
        self.table.clearContents()
        self.table.setRowCount(0)  # 清除所有行



def main():
    app = QApplication(sys.argv)
    window = MyApp()  # 实例化你的窗口类
    window.show()  # 显示窗口
    print("启动")
    sys.exit(app.exec_())  # 进入主事件循环，并在退出时确保应用正确关闭


if __name__ == '__main__':
    main()